self

\(Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V\)

• \(X\): context-free embedding of the input sequence

  • \(X \in R^{n \times d_{model}}\)
  • \(n\): sequence length
  • \(d_{model}\): embedding dimension

• \(Q = X W_Q\)

• \(K = X W_K\)

• \(V = X W_V\)

  • where \(W_Q, W_K, W_V\) are learnable parameters, called projection matrices
  • \(W_Q \in R^{d_{model} \times d_k}\): what to look for
  • \(W_K \in R^{d_{model} \times d_k}\): what to compare with
  • \(W_V \in R^{d_{model} \times d_v}\): what information to extract

• Attention:

  • score: \(S = QK^T\)
  • scaled: \(S_{scaled} = \frac{S}{\sqrt{d_k}}\)
  • softmax: \(A = softmax(S_{scaled})\)

• Final output: \(Attention(Q, K, V) = AV\)

  • the whole process lifts the context-free embedding into a contextualized representation

• \(d_k\) and \(d_v\) are hyperparameters

  • often \(d_k = d_v = d_{model} / h\)
  • \(h\): number of attention heads

Causal Masking:

• mask = [ [0, -∞, -∞], [0, 0, -∞], [0, 0, 0] ]
• \(A = softmax(S_{scaled} + mask)\)
• In decoder self-attention, we're generating from left to right. During training, we have the full target sequence available, but we don't want to attend to future positions.

Scalability:

• Time complexity: \(O(n^2 d_{model})\)
• Space complexity: \(O(n^2)\)
  • for storing the attention matrix (\(S_{scaled}\))
• Quadratic scaling in context length is the reason why transformer struggles in long context

Multi-head:

• \(MultiHead(Q, K, V) = Concat(head_1, head_2, ..., head_h) W^O\)
• \(head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)\)

Parameter calculation:

• \(W_Q^i, W_K^i \in R^{d_{model} \times d_k \times h}\) for each head \(i\)
• \(W_V^i \in R^{d_{model} \times d_v \times h}\) for each head \(i\)
• \(W^O \in R^{h \times d_v \times d_{model}}\)